package com.stockprediction.analysis;

import java.util.List;
import java.util.Scanner;
import java.util.stream.Collectors;
import libsvm.*;

public class                                                                                                                                                                                                                                                                                  Main {
    public static void main(String[] args) throws Exception {
        System.out.println("📥 读取训练集...");
        List<NewsData> trainData = CSVReader.readCSV("C:\\Users\\Xander.Zheng\\Desktop\\train_data.csv");
        System.out.println("训练集大小: " + trainData.size());
        System.out.println("📥 读取测试集...");
        List<NewsData> testData = CSVReader.readCSV("C:\\Users\\Xander.Zheng\\Desktop\\test_data.csv");

        System.out.println("🔍 训练集分词...");
        List<String> trainTexts = trainData.stream()
                .map(news -> news.getTitle() + " " + news.getContent()) // 拼接标题+正文
                .map(ChineseTokenizer::tokenize)
                .collect(Collectors.toList());
        List<Integer> trainLabels = trainData.stream().map(NewsData::getLabel).toList();

        System.out.println("📊 计算 TF-IDF...");
        TFIDFVectorizer vectorizer = new TFIDFVectorizer();
        vectorizer.fit(trainTexts);
        // **保存 TF-IDF 向量化器**
        vectorizer.save("tfidf_vectorizer.bin");
        List<double[]> X_train = trainTexts.stream().map(vectorizer::transform).toList();

        System.out.println("🚀 训练 SVM...");
        svm_model model = SVMTrainer.train(X_train, trainLabels);

        System.out.println("🔍 测试集分词...");
        List<String> testTexts = testData.stream()
                .map(news -> news.getTitle() + " " + news.getContent()) // 拼接标题+正文
                .map(ChineseTokenizer::tokenize)
                .collect(Collectors.toList());
        List<Integer> testLabels = testData.stream().map(NewsData::getLabel).toList();

        System.out.println("📊 转换测试数据 TF-IDF...");
        List<double[]> X_test = testTexts.stream().map(vectorizer::transform).toList();

        System.out.println("🧐 评估模型...");
        evaluateModel(model, X_test, testLabels);

        System.out.println("✅ 训练和测试完成！");

        // 🏆 手动输入新闻测试
        predictSingleNews(vectorizer, model);
    }

    // 评估模型
    public static void evaluateModel(svm_model model, List<double[]> X_test, List<Integer> y_test) {
        int correct = 0;
        for (int i = 0; i < X_test.size(); i++) {
            int predicted = predict(model, X_test.get(i));
            if (predicted == y_test.get(i)) {
                correct++;
            }
        }
        double accuracy = (double) correct / y_test.size();
        System.out.println("🎯 测试集准确率: " + String.format("%.4f", accuracy));
    }

    // 进行单条新闻预测
    public static int predict(svm_model model, double[] features) {
        svm_node[] nodes = new svm_node[features.length];
        for (int i = 0; i < features.length; i++) {
            nodes[i] = new svm_node();
            nodes[i].index = i + 1;
            nodes[i].value = features[i];
        }
        return (int) svm.svm_predict(model, nodes);
    }



    // 🎯 手动输入新闻进行预测
    public static void predictSingleNews(TFIDFVectorizer vectorizer, svm_model model) {
        Scanner scanner = new Scanner(System.in);
        while (true) {
            System.out.println("\n🔹 请输入一条新闻 (输入 'exit' 退出): ");
            String inputText = scanner.nextLine();
            if (inputText.equalsIgnoreCase("exit")) {
                System.out.println("👋 退出测试！");
                break;
            }

            // 1️⃣ 分词
            String tokenizedText = ChineseTokenizer.tokenize(inputText);

            // 2️⃣ 转换为 TF-IDF 特征
            double[] features = vectorizer.transform(tokenizedText);

            // 3️⃣ 进行 SVM 预测
            int prediction = predict(model, features);

            // 4️⃣ 输出预测结果
            String sentiment = (prediction == 1) ? "😃 正面" : "😡 负面";
            System.out.println("📢 预测情感分类：" + sentiment);
        }
        scanner.close();
    }
}
